import sys
import cv2
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QPushButton, QVBoxLayout, QFileDialog
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt
from ultralytics import YOLO

class PersonDetector(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("YOLO 人体检测")
        self.model = YOLO("yolov8n.pt")  # 使用 yolov8n 模型，适合快速演示
        self.init_ui()

    def init_ui(self):
        self.label = QLabel("请选择一张图像")
        self.label.setAlignment(Qt.AlignCenter)

        self.button = QPushButton("加载图像并检测")
        self.button.clicked.connect(self.load_and_detect)

        layout = QVBoxLayout()
        layout.addWidget(self.label)
        layout.addWidget(self.button)
        self.setLayout(layout)

    def load_and_detect(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图像", "", "Images (*.png *.jpg *.jpeg)")
        if file_path:
            img = cv2.imread(file_path)
            results = self.model(img)[0]

            for box in results.boxes:
                cls_id = int(box.cls[0])
                if cls_id == 0:  # 只画 "person"
                    x1, y1, x2, y2 = map(int, box.xyxy[0])
                    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                    cv2.putText(img, 'Person', (x1, y1 - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)

            # 转换为 Qt 显示的格式
            rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_img.shape
            bytes_per_line = ch * w
            qt_img = QImage(rgb_img.data, w, h, bytes_per_line, QImage.Format_RGB888)
            pixmap = QPixmap.fromImage(qt_img)
            self.label.setPixmap(pixmap.scaled(self.label.size(), Qt.KeepAspectRatio))

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = PersonDetector()
    window.resize(800, 600)
    window.show()
    sys.exit(app.exec_())
